#!/usr/bin/python3
import os
import cv2
import numpy as np
import struct

person_select_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/FaceRecSel.txt'
video_label_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/video_label_list.txt'

feat_dim = 512


def read_features(file_list, feature_file, feat_dim=512, skip_offset=0):
    video_image_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 1000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/sqqjd/ycjy/150100414a54443452064d2742d03600/video/20200509_173513/00110_000.jpg
            lines = line.strip().split()
            file_path = lines[0].strip()

            feature = feat_file.read(feat_dim * 4 + skip_offset)
            feature = feature[skip_offset:]

            video_name = os.path.dirname(file_path)
            file_name = os.path.basename(file_path)
            image_features = video_image_features.get(video_name, None)
            if image_features is None:
                video_image_features[video_name] = [[file_name, feature]]
            else:
                image_features.append([file_name, feature])

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return video_image_features


def main():
    video_genders = {}
    with open(video_label_list_file, 'r') as file:
        for line in file.readlines():
            # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_154308 1
            lines = line.split()
            video_name = lines[0].strip()
            gender = lines[1].strip()
            video_genders[video_name] = int(gender)

    video_images = {}
    faceid_genders = {}
    faceid_videos = {}
    with open(person_select_list_file, 'r') as file:
        for line in file.readlines():
            # 1687 1687 gd/zs/nqjd/dxjyjy/150100414a54443452064d2742e83600/video/20190702_174057/00131_000.jpg 0
            lines = line.split()
            faceid = '{} {}'.format(lines[0].strip(), lines[1].strip())
            video_name = os.path.dirname(lines[2])
            image_name = os.path.basename(lines[2])
            gender = video_genders.get(video_name, None)
            if gender is None:
                continue

            genders = faceid_genders.get(faceid, None)
            if genders is None:
                faceid_genders[faceid] = {gender}
            else:
                genders.add(gender)

            videos = faceid_videos.get(faceid, None)
            if videos is None:
                faceid_videos[faceid] = {video_name}
            else:
                videos.add(video_name)

            images = video_images.get(video_name, None)
            if images is None:
                video_images[video_name] = [image_name]
            else:
                images.append(image_name)

    out_path = os.path.dirname(person_select_list_file)

    out_file_path = os.path.join(out_path, 'person_image_gender.txt')
    out_file = open(out_file_path, 'w')

    out_faceid_file_path = os.path.join(out_path, 'faceid_person_image_gender.txt')
    out_faceid_file = open(out_faceid_file_path, 'w')

    video_image_genders = []
    for faceid, genders in faceid_genders.items():
        if len(genders) > 1:
            print(faceid, genders)
            continue
        videos = faceid_videos[faceid]
        genders = list(genders)
        for video in videos:
            images = video_images[video]
            video_image_genders.append([video, faceid, images, genders[0]])

    video_image_genders.sort(key=lambda x: x[0])
    video_count = 0
    file_count = 0
    for video, faceid, images, gender in video_image_genders:
        video_count += 1
        images.sort()

        for image in images:
            out_info = '{}/{} {}\n'.format(video, image, gender)
            out_file.write(out_info)
            out_info = '{} {}/{} {}\n'.format(faceid, video, image, gender)
            out_faceid_file.write(out_info)
            file_count += 1
    out_file.close()
    out_faceid_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {}'.format(out_faceid_file_path))
    print('Video Count: {}, File Count: {}'.format(video_count, file_count))


if __name__ == '__main__':
    main()
